package com.jd.genie.agent.agent;

import com.jd.genie.agent.dto.Message;
import com.jd.genie.agent.dto.tool.ToolCall;
import com.jd.genie.agent.dto.tool.ToolChoice;
import com.jd.genie.agent.enums.AgentState;
import com.jd.genie.agent.enums.RoleType;
import com.jd.genie.agent.llm.LLM;
import com.jd.genie.agent.printer.Printer;
import com.jd.genie.agent.tool.BaseTool;
import com.jd.genie.agent.tool.common.PlanningTool;
import com.jd.genie.agent.util.FileUtil;
import com.jd.genie.agent.util.SpringContextHolder;
import com.jd.genie.config.GenieConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationContext;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;

/**
 * 规划代理 - 创建和管理任务计划的代理
 */
@Slf4j
public class PlanningAgent extends ReActAgent {

    private List<ToolCall> toolCalls;
    private PlanningTool planningTool = new PlanningTool();
    private boolean isColseUpdate;
    private String systemPromptSnapshot;
    private String nextStepPromptSnapshot;

    /**
     * {@link GenieConfig#getGenieSopPrompt()}
     */
    private final String plannerSystemPrompt;

    {
        final String SYSTEM_PROMPT = """
                {{sopPrompt}}
                
                ===
                # 环境变量
                ## 当前日期
                {{date}}
                
                # 当前可用的文件名及描述
                {{files}}
                
                # 约束
                - 思考过程中，不要透露你的工具名称
                - 调用planning生成任务列表，完成所有子任务就能完成任务。
                - 以上是你需要遵循的指令。
                
                Let's think step by step (让我们一步步思考)
                """;
        plannerSystemPrompt = """
                
                # 角色
                你是一个智能助手，名叫Genie。
                
                # 说明
                你是任务规划助手，根据用户需求，拆解任务列表，从而确定planning工具入参。每次执行planning工具前，必须先输出本轮思考过程（reasoning），再调用planning工具生成任务列表。
                
                # 技能
                - 擅长将用户任务拆解为具体、独立的任务列表。
                - 对简单任务，避免过度拆解任务。
                - 对复杂任务，合理拆解为多个有逻辑关联的子任务
                
                # 处理需求
                ## 拆解任务
                - 深度推理分析用户输入，识别核心需求及潜在挑战。
                - 将复杂问题分解为可管理、可执行、独立且清晰的子任务，任务之间不重复、不交叠。拆解最多不超过5个任务。
                - 任务按顺序或因果逻辑组织，上下任务逻辑连贯。
                - 读取文件后，对文件进行处理，处理完成保存文件应该放到一个子任务中。
                
                ## 要求
                - 每一个子任务都是一个完整的子任务，例如读取文件后，将文件中的表格抽取出出来形成表格保存。
                - 调用planning工具前，必须输出500字以内的思考过程，说明本轮任务拆解的依据与目标。
                - 首次规划拆分时，输出整体拆分思路；后续如需调整，也需输出调整思考。
                - 每个子任务为清晰、独立的指令，细化完成标准，不重复、不交叠。
                - 不要输出重复的任务。
                - 任务中间不能输出网页版报告，只能在最后一个任务中，生成一个网页版报告。
                - 最后一个任务是需要输出报告时，如果没有明确要求，优先“输出网页版报告”，如果有指定格式要求，最后一个任务按用户指定的格式输出。
                - 当前不能支持用户在计划中提供内容，因此不要要求用户提供信息
                
                ## 输出格式
                输出本轮思考过程，200字以内，简明说明拆解任务依据或调整依据，并调用planning工具生成任务计划。
                
                # 语言设置
                - 所有内容均以 **中文** 输出
                
                # 任务示例：
                以下仅是你拆解任务的一个简单参考示例，你在解决问题时，参考如下拆解任务，但不要局限于如下示例计划
                
                ## 示例任务1：分析 xxxx
                任务列表
                - 执行顺序1. 信息收集：收集xxxx
                - 执行顺序2. 筛选分析：xxxx，分析并保存成Markdown文件
                - 执行顺序3. 输出报告：以网页形式呈现分析报告，调用网页生成工具
                
                ##示例任务2：提取文件中的表格
                任务列表
                - 执行顺序1. 文件表格提取：读取文件内容，抽取文件中存在的表格，并保存成表格文件。
                
                ## 示例任务3：分析 xxxx，以PPT格式展示
                任务列表
                - 执行顺序1. 信息收集：收集xxxx
                - 执行顺序2. 筛选分析：xxxx，分析并保存成Markdown文件
                - 执行顺序3. 输出PPT：以PPT呈现xx，调用PPT生成工具
                
                ## 示例任务4：我要写一个 xxxx
                任务列表
                - 执行顺序1. 信息收集：收集xxxx
                - 执行顺序2. 文件输出：以网页形式呈现xxx，调用网页生成工具
                
                
                ===
                # 环境变量
                ## 当前日期
                <date>
                {{date}}
                </date>
                
                ## 当前可用的文件名及描述
                <files>
                {{files}}
                </files>
                
                ## 用户历史对话信息
                <history_dialogue>
                {{history_dialogue}}
                </history_dialogue>
                
                ## 约束
                - 思考过程中，不要透露你的工具名称
                - 调用planning生成任务列表，完成所有子任务就能完成任务。
                - 以上是你需要遵循的指令，不要输出在结果中。
                
                Let's think step by step (让我们一步步思考)
                
                """;

    }

    private final String plannerNextStepPrompt;

    {
        final String NEXT_STEP_PROMPT = """
                工具planing的参数有
                必填参数1：命令command
                可选参数2：当前步状态step_status。
                
                必填参数1：命令command的枚举值有：
                'mark_step', 'finish'
                含义如下：
                - 'finish' 根据已有的执行结果，可以判断出任务已经完成，输出任务结束，命令command为：finish
                - 'mark_step' 标记当前任务规划的状态，设置当前任务的step_status
                
                当参数command值为mark_step时，需要可选参数2step_status，其中当前步状态step_status的枚举值如下：
                - 没有开始'not_started'
                - 进行中'in_progress'
                - 已完成'completed'
                
                对应如下几种情况：
                1.当前任务是否执行完成，完成以及失败都算执行完成，执行完成将入参step_status设置为`completed`
                
                一步一步分析完成任务，确定工具planing的入参，调用planing工具
                """;
        plannerNextStepPrompt = """
                工具planing的参数有
                必填参数1：命令command
                可选参数2：当前步状态step_status。
                
                必填参数1：命令command的枚举值有：
                'mark_step', 'finish'
                含义如下：
                - 'finish' 根据已有的执行结果，可以判断出任务已经完成，输出任务结束，命令command为：finish
                - 'mark_step' 标记当前任务规划的状态，设置当前任务的step_status
                
                当参数command值为mark_step时，需要可选参数2step_status，其中当前步状态step_status的枚举值如下：
                - 没有开始'not_started'
                - 进行中'in_progress'\s
                - 已完成'completed'
                
                对应如下几种情况：
                1.当前任务是否执行完成，完成以及失败都算执行完成，执行完成将入参step_status设置为`completed`
                
                一步一步分析完成任务，确定工具planing的入参，调用planing工具
                """;
    }

    // An agent that creates and manages plans to solve tasks
    public PlanningAgent(AgentContext context, Printer printer)
    {
        super(context);

        ApplicationContext applicationContext = SpringContextHolder.getApplicationContext();
        GenieConfig genieConfig = applicationContext.getBean(GenieConfig.class);

        StringBuilder toolPrompt = new StringBuilder();
        for (BaseTool tool : context.getToolCollection().getToolMap().values()) {
            toolPrompt.append(String.format("工具名：%s 工具描述：%s\n", tool.getName(), tool.getDescription()));
        }

        setSystemPrompt(plannerSystemPrompt
                .replace("{{tools}}", toolPrompt.toString())
                .replace("{{query}}", context.getQuery())
                .replace("{{date}}", context.getDateInfo()));
        setNextStepPrompt(plannerNextStepPrompt
                .replace("{{tools}}", toolPrompt.toString())
                .replace("{{query}}", context.getQuery())
                .replace("{{date}}", context.getDateInfo()));

        systemPromptSnapshot = getSystemPrompt();
        nextStepPromptSnapshot = getNextStepPrompt();

        setPrinter(printer);
        setLlm(new LLM(genieConfig.getPlannerModelName(), ""));

        isColseUpdate = "1".equals(genieConfig.getPlanningCloseUpdate());

        // 初始化工具集合
        availableTools.addTool(planningTool);
        planningTool.setAgentContext(context);
    }

    public boolean think()
    {
        // 获取文件内容
        String filesStr = FileUtil.formatFileInfo(context.getProductFiles(), false);
        setSystemPrompt(systemPromptSnapshot.replace("{{files}}", filesStr));
        setNextStepPrompt(nextStepPromptSnapshot.replace("{{files}}", filesStr));
        log.info("{} planer fileStr {}", context.getRequestId(), filesStr);

        // 关闭了动态更新Plan，直接执行下一个task
        if (isColseUpdate) {
            if (Objects.nonNull(planningTool.getPlan())) {
                planningTool.stepPlan();
                return true;
            }
        }

        try {
            if (!getMemory().getLastMessage().getRole().equals(RoleType.USER)) {
                Message userMsg = Message.userMessage(getNextStepPrompt(), null);
                getMemory().addMessage(userMsg);
            }

            context.setStreamMessageType("plan_thought");
            LLM.ToolCallResponse response = getLlm().askTool(context,
                    getMemory().getMessages(),
                    Message.systemMessage(getSystemPrompt(), null),
                    availableTools,
                    ToolChoice.AUTO, null, context.getIsStream(), 300,
                    printer);

            toolCalls = response.getToolCalls();

            // 记录响应信息
            if (!context.getIsStream() && response.getContent() != null && !response.getContent().isEmpty()) {
                printer.send("plan_thought", response.getContent());
            }

            // 记录响应信息
            String name = getClass().getSimpleName();
            log.info("{} {}'s thoughts: {}", context.getRequestId(), name, response.getContent());
            log.info("{} {} selected {} tools to use", context.getRequestId(), name, response.getToolCalls() != null ? response.getToolCalls().size() : 0);

            // 创建并添加助手消息
            Message assistantMsg = response.getToolCalls() != null && !response.getToolCalls().isEmpty() ?
                    Message.fromToolCalls(response.getContent(), response.getToolCalls()) :
                    Message.assistantMessage(response.getContent(), null);

            getMemory().addMessage(assistantMsg);
        } catch (Exception e) {
            log.error("{} think error ", context.getRequestId(), e);
        }
        return true;
    }

    public String act()
    {
        // 关闭了动态更新Plan，直接执行下一个task
        if (isColseUpdate) {
            if (Objects.nonNull(planningTool.getPlan())) {
                return getNextTask();
            }
        }
        Integer maxObserve = null;
        List<String> results = new ArrayList<>();
        for (ToolCall toolCall : toolCalls) {
            String result = executeTool(toolCall);
            if (maxObserve != null) {
                result = result.substring(0, Math.min(result.length(), maxObserve));
            }
            results.add(result);

            // 添加工具响应到记忆
            // function_call
            Message toolMsg = Message.toolMessage(
                    result,
                    toolCall.getId(),
                    null
            );
            getMemory().addMessage(toolMsg);
        }

        if (Objects.nonNull(planningTool.getPlan())) {
            if (isColseUpdate) {
                planningTool.stepPlan();
            }
            return getNextTask();
        }

        return String.join("\n\n", results);
    }


    private String getNextTask()
    {
        boolean allComplete = true;
        for (String status : planningTool.getPlan().getStepStatus()) {
            if (!"completed".equals(status)) {
                allComplete = false;
                break;
            }
        }

        if (allComplete) {
            setState(AgentState.FINISHED);
            printer.send("plan", planningTool.getPlan());
            return "finish";
        }

        if (!planningTool.getPlan().getCurrentStep().isEmpty()) {
            setState(AgentState.FINISHED);
            String[] currentSteps = planningTool.getPlan().getCurrentStep().split("<sep>");
            printer.send("plan", planningTool.getPlan());
            Arrays.stream(currentSteps).forEach(step -> printer.send("task", step));
            return planningTool.getPlan().getCurrentStep();
        }
        return "";
    }

    private final String planPrePrompt = "一步一步（step by step）思考，结合用户上传的文件分析用户问题，并根据问题制定计划，用户问题如下：";

    public String run(String request)
    {
        if (Objects.isNull(planningTool.getPlan())) {
            request = planPrePrompt + request;
        }
        return baseRun(request, this::think, this::act);
    }

}